"""
负载测试脚本
使用Locust进行系统负载测试
"""

from locust import HttpUser, task, between
import json
import random
import time
from typing import Dict, Any


class CampusVideoTwinUser(HttpUser):
    """校园视频孪生系统用户行为模拟"""
    
    wait_time = between(1, 3)  # 用户操作间隔1-3秒
    
    def on_start(self):
        """用户开始时的初始化操作"""
        self.camera_ids = []
        self.alert_ids = []
        self.setup_test_data()
    
    def setup_test_data(self):
        """设置测试数据"""
        # 创建测试摄像头
        for i in range(3):
            camera_data = {
                "name": f"负载测试摄像头{i}",
                "location": f"测试位置{i}",
                "ip_address": f"192.168.1.{150+i}",
                "port": 554,
                "username": "admin",
                "password": "123456",
                "stream_url": f"rtsp://192.168.1.{150+i}:554/stream1",
                "status": "online",
                "position": {"x": float(i*100), "y": float(i*50), "z": 10.0}
            }
            
            response = self.client.post("/api/v1/cameras/", json=camera_data)
            if response.status_code == 200:
                camera = response.json()
                self.camera_ids.append(camera["id"])
    
    @task(10)
    def get_cameras(self):
        """获取摄像头列表 - 高频操作"""
        with self.client.get("/api/v1/cameras/", catch_response=True) as response:
            if response.status_code == 200:
                cameras = response.json()
                response.success()
            else:
                response.failure(f"获取摄像头列表失败: {response.status_code}")
    
    @task(8)
    def get_system_health(self):
        """获取系统健康状态 - 高频操作"""
        with self.client.get("/api/v1/system/health", catch_response=True) as response:
            if response.status_code == 200:
                health = response.json()
                if health.get("status") == "healthy":
                    response.success()
                else:
                    response.failure("系统状态不健康")
            else:
                response.failure(f"获取系统健康状态失败: {response.status_code}")
    
    @task(6)
    def get_alerts(self):
        """获取告警列表 - 中频操作"""
        with self.client.get("/api/v1/alerts/", catch_response=True) as response:
            if response.status_code == 200:
                alerts = response.json()
                response.success()
            else:
                response.failure(f"获取告警列表失败: {response.status_code}")
    
    @task(5)
    def get_camera_detail(self):
        """获取摄像头详情 - 中频操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        with self.client.get(f"/api/v1/cameras/{camera_id}", catch_response=True) as response:
            if response.status_code == 200:
                camera = response.json()
                response.success()
            else:
                response.failure(f"获取摄像头详情失败: {response.status_code}")
    
    @task(4)
    def get_video_stream_url(self):
        """获取视频流地址 - 中频操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        with self.client.get(f"/api/v1/video-stream/{camera_id}/url", catch_response=True) as response:
            if response.status_code == 200:
                stream_info = response.json()
                response.success()
            else:
                response.failure(f"获取视频流地址失败: {response.status_code}")
    
    @task(3)
    def start_ai_detection(self):
        """启动AI检测任务 - 中频操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        task_data = {
            "camera_id": camera_id,
            "task_type": "person_detection",
            "parameters": {
                "confidence_threshold": 0.5,
                "detection_interval": 1.0
            }
        }
        
        with self.client.post("/api/v1/ai-algorithm/person-detection/start", 
                            json=task_data, catch_response=True) as response:
            if response.status_code == 200:
                result = response.json()
                response.success()
            else:
                response.failure(f"启动AI检测任务失败: {response.status_code}")
    
    @task(3)
    def get_detection_results(self):
        """获取检测结果 - 中频操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        with self.client.get(f"/api/v1/ai-algorithm/person-detection/results/{camera_id}", 
                           catch_response=True) as response:
            if response.status_code == 200:
                results = response.json()
                response.success()
            else:
                response.failure(f"获取检测结果失败: {response.status_code}")
    
    @task(2)
    def create_alert(self):
        """创建告警 - 低频操作"""
        alert_data = {
            "title": f"负载测试告警{random.randint(1, 1000)}",
            "message": "这是一个负载测试生成的告警",
            "level": random.choice(["low", "medium", "high", "critical"]),
            "location": f"测试位置{random.randint(1, 10)}",
            "alert_type": "load_test",
            "metadata": {"load_test": True, "timestamp": time.time()}
        }
        
        with self.client.post("/api/v1/alerts/", json=alert_data, catch_response=True) as response:
            if response.status_code == 200:
                alert = response.json()
                self.alert_ids.append(alert["id"])
                response.success()
            else:
                response.failure(f"创建告警失败: {response.status_code}")
    
    @task(2)
    def acknowledge_alert(self):
        """确认告警 - 低频操作"""
        if not self.alert_ids:
            return
        
        alert_id = random.choice(self.alert_ids)
        with self.client.post(f"/api/v1/alerts/{alert_id}/acknowledge", 
                            catch_response=True) as response:
            if response.status_code == 200:
                result = response.json()
                response.success()
            else:
                response.failure(f"确认告警失败: {response.status_code}")
    
    @task(1)
    def get_system_statistics(self):
        """获取系统统计 - 低频操作"""
        with self.client.get("/api/v1/system/statistics", catch_response=True) as response:
            if response.status_code == 200:
                stats = response.json()
                response.success()
            else:
                response.failure(f"获取系统统计失败: {response.status_code}")
    
    @task(1)
    def analyze_crowd_density(self):
        """分析人流密度 - 低频操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        analysis_data = {
            "camera_id": camera_id,
            "task_type": "crowd_density",
            "parameters": {
                "analysis_interval": 5.0,
                "density_threshold": 0.8
            }
        }
        
        with self.client.post("/api/v1/data-analysis/crowd-density/analyze", 
                            json=analysis_data, catch_response=True) as response:
            if response.status_code == 200:
                result = response.json()
                response.success()
            else:
                response.failure(f"分析人流密度失败: {response.status_code}")


class AdminUser(HttpUser):
    """管理员用户行为模拟"""
    
    wait_time = between(2, 5)  # 管理员操作间隔较长
    weight = 1  # 管理员用户权重较低
    
    def on_start(self):
        """管理员用户初始化"""
        self.camera_ids = []
    
    @task(5)
    def create_camera(self):
        """创建摄像头 - 管理员操作"""
        camera_data = {
            "name": f"管理员创建摄像头{random.randint(1, 1000)}",
            "location": f"管理员测试位置{random.randint(1, 20)}",
            "ip_address": f"192.168.2.{random.randint(100, 200)}",
            "port": 554,
            "username": "admin",
            "password": "123456",
            "stream_url": f"rtsp://192.168.2.{random.randint(100, 200)}:554/stream1",
            "status": "online",
            "position": {
                "x": float(random.randint(0, 1000)),
                "y": float(random.randint(0, 1000)),
                "z": 10.0
            }
        }
        
        with self.client.post("/api/v1/cameras/", json=camera_data, catch_response=True) as response:
            if response.status_code == 200:
                camera = response.json()
                self.camera_ids.append(camera["id"])
                response.success()
            else:
                response.failure(f"创建摄像头失败: {response.status_code}")
    
    @task(3)
    def update_camera(self):
        """更新摄像头 - 管理员操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        update_data = {
            "name": f"更新后的摄像头{random.randint(1, 1000)}",
            "location": f"更新后的位置{random.randint(1, 20)}"
        }
        
        with self.client.put(f"/api/v1/cameras/{camera_id}", 
                           json=update_data, catch_response=True) as response:
            if response.status_code == 200:
                camera = response.json()
                response.success()
            else:
                response.failure(f"更新摄像头失败: {response.status_code}")
    
    @task(2)
    def test_camera_connection(self):
        """测试摄像头连接 - 管理员操作"""
        if not self.camera_ids:
            return
        
        camera_id = random.choice(self.camera_ids)
        with self.client.post(f"/api/v1/cameras/{camera_id}/test-connection", 
                            catch_response=True) as response:
            if response.status_code == 200:
                result = response.json()
                response.success()
            else:
                response.failure(f"测试摄像头连接失败: {response.status_code}")
    
    @task(1)
    def delete_camera(self):
        """删除摄像头 - 管理员操作"""
        if not self.camera_ids:
            return
        
        camera_id = self.camera_ids.pop()
        with self.client.delete(f"/api/v1/cameras/{camera_id}", catch_response=True) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"删除摄像头失败: {response.status_code}")


class MonitorUser(HttpUser):
    """监控用户行为模拟"""
    
    wait_time = between(0.5, 2)  # 监控用户操作频率最高
    weight = 3  # 监控用户权重最高
    
    @task(15)
    def monitor_dashboard(self):
        """监控仪表板 - 高频刷新"""
        endpoints = [
            "/api/v1/system/health",
            "/api/v1/alerts/",
            "/api/v1/cameras/",
            "/api/v1/system/statistics"
        ]
        
        endpoint = random.choice(endpoints)
        with self.client.get(endpoint, catch_response=True) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"监控仪表板请求失败: {response.status_code}")
    
    @task(5)
    def check_alert_status(self):
        """检查告警状态"""
        with self.client.get("/api/v1/alerts/?status=active", catch_response=True) as response:
            if response.status_code == 200:
                alerts = response.json()
                response.success()
            else:
                response.failure(f"检查告警状态失败: {response.status_code}")


# 自定义负载测试场景
class StressTestUser(HttpUser):
    """压力测试用户 - 模拟极端负载"""
    
    wait_time = between(0.1, 0.5)  # 极短间隔
    weight = 1
    
    @task
    def stress_test_endpoints(self):
        """压力测试各个端点"""
        endpoints = [
            "/api/v1/system/health",
            "/api/v1/cameras/",
            "/api/v1/alerts/"
        ]
        
        # 快速连续请求
        for endpoint in endpoints:
            with self.client.get(endpoint, catch_response=True) as response:
                if response.status_code != 200:
                    response.failure(f"压力测试失败: {endpoint} - {response.status_code}")
                else:
                    response.success()


if __name__ == "__main__":
    # 可以直接运行此脚本进行本地测试
    import subprocess
    import sys
    
    print("启动负载测试...")
    print("使用命令: locust -f load_test.py --host=http://localhost:8000")
    
    # 自动启动locust
    try:
        subprocess.run([
            "locust", 
            "-f", __file__, 
            "--host=http://localhost:8000",
            "--users=10",
            "--spawn-rate=2",
            "--run-time=60s",
            "--headless"
        ])
    except FileNotFoundError:
        print("请先安装locust: pip install locust")
        sys.exit(1)